package com.kingsoft.dc.khaos.module.spark.preprocess.transform

import com.kingsoft.dc.khaos.KhaosContext
import com.kingsoft.dc.khaos.innertype.Schema
import com.kingsoft.dc.khaos.metadata.{Dependency, KhaosStructField}
import com.kingsoft.dc.khaos.module.spark.constants.ColumnType
import com.kingsoft.dc.khaos.module.spark.metadata.preprocess.transform.{AggAdvancedOption, AggOptionInfo, AggregateInfo}
import com.kingsoft.dc.khaos.module.spark.util.DataframeUtils
import org.apache.spark.sql.{Column, DataFrame}
import org.json4s.DefaultFormats
import org.json4s.jackson.JsonMethods.parse
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.DecimalType

/**
  * Created by haorenhui on 2019/07/17.
  */
class Aggregate  extends TransformStrategy {

    //连接方式枚举
    object AggregateEnum {
        val min = "min"
        val max = "max"
        val sum = "sum"
        val avg = "avg"
        val count = "count"
        val countdistinct = "countdistinct"
        val empty = ""
    }

    private var aggregateInfo: AggregateInfo = _

    override def exec(kc: KhaosContext,
                      module_id: String,
                      config: String,
                      dependences: Seq[Dependency],
                      targets:Seq[Dependency]): Seq[(String,DataFrame)] = {
        //获取配置信息
        implicit val formats:DefaultFormats = DefaultFormats
        aggregateInfo = parse(config, true:Boolean).extract[AggregateInfo]

        //获取上游信息
        var (schema, dataFrame) = kc.structData[DataFrame](dependences.head.getEdge())
        val group_fields: List[String] = aggregateInfo.group_fields
        val extract_fields: List[AggOptionInfo] = aggregateInfo.extract_fields.filter(field=>field.option.nonEmpty)
        val groupArr: List[Column] = group_fields.map(col(_))

        extract_fields.foreach(tp=>{
            val field: String = tp.field
            val data_type: String = tp.data_type
            if(data_type.equals(ColumnType.DECIMAL)){
                val length: String = tp.length.get
                if(length.contains(",")){
                    val realLen: Int = length.split("\\,")(0).toInt
                    val decimals: Int = length.split("\\,")(1).toInt
                    dataFrame = dataFrame.withColumn(field,dataFrame.col(field).cast(DecimalType(realLen,decimals)))
                }
            }
        })

        //rename 添加字段后缀
        val aggArr: Array[Column] = extract_fields.map(col => {
            val field: String = col.field
            var result: Column = null
            col.option match {
                case AggregateEnum.min => result = min(field) as s"${field}_min"
                case AggregateEnum.max => result = max(field) as s"${field}_max"
                case AggregateEnum.sum => result = sum(field) as s"${field}_sum"
                case AggregateEnum.avg => result = avg(field) as s"${field}_avg"
                case AggregateEnum.count => result = count(field) as s"${field}_count"
                case AggregateEnum.countdistinct => result = countDistinct(field) as s"${field}_countdistinct"
                case _ => throw new Exception(s"不支持的聚合类型${col.option}")
            }
            result
        }).toArray

        // 聚合
        var result:DataFrame = dataFrame.groupBy(groupArr:_*).agg(aggArr(0),aggArr.slice(1,aggArr.length):_*)

        //过滤
        val filter: List[String] = aggregateInfo.filter
        if(filter.nonEmpty){
            filter.foreach(filt=>
                if(filt.nonEmpty)
                    result=result.where(filt)
            )
        }

        //分区 缓存
        val advanced_options: AggAdvancedOption = aggregateInfo.advanced_options
        if (true.equals(advanced_options.cache_option.on_off)) {
            result = DataframeUtils.cacheDataframe(result, this)
        }
        if (true.equals(advanced_options.partition.on_off)) {
            result = DataframeUtils.repartionDataframe(result, aggregateInfo.advanced_options.partition.nums.toInt, this)
        }

        addResult(targets.head,result)
    }

    override def schema(kc: KhaosContext,
                        module_id: String,
                        config: String,
                        dependencies: Seq[Dependency]) :Schema= {
        //上游DF schema
        val headSchema:Schema = kc.schemaChannel.getSchema(dependencies.head.getSource())
        //(字段,数据类型)
        val headSchemaMap: Map[String, String] = headSchema.toList().map(schema=>(schema.name,schema.typeTag)).toMap
        //生成聚合字段schema
        var schemaMap: Map[String, String] = aggregateInfo.extract_fields.map(aggInfo => {
            var fieldName: String = aggInfo.field
            var fieldDataType: String = ""
            aggInfo.option match {
                case AggregateEnum.min => fieldDataType = headSchemaMap.getOrElse(fieldName,ColumnType.STRING); fieldName=fieldName + "_" +  aggInfo.option
                case AggregateEnum.max => fieldDataType = headSchemaMap.getOrElse(fieldName,ColumnType.STRING); fieldName=fieldName + "_" +  aggInfo.option
                case AggregateEnum.sum => fieldDataType = ColumnType.DECIMAL; fieldName=fieldName + "_" +  aggInfo.option
                case AggregateEnum.avg => fieldDataType = ColumnType.DECIMAL; fieldName=fieldName + "_" +  aggInfo.option
                case AggregateEnum.count => fieldDataType = ColumnType.NUMBER; fieldName=fieldName + "_" +  aggInfo.option;
                case AggregateEnum.countdistinct => fieldDataType = ColumnType.NUMBER; fieldName=fieldName + "_" +  aggInfo.option;
                case _ => fieldDataType = ColumnType.STRING; fieldName = fieldName
            }
            (fieldName, fieldDataType)
        }).toMap
        //生成分组字段schema
        val groupFieldsMap:Map[String,String] = aggregateInfo.group_fields.map(groupFieldName => {
            val data_type: String = headSchemaMap.getOrElse(groupFieldName,ColumnType.STRING)
            (groupFieldName, data_type)
        }).toMap
        schemaMap=schemaMap ++ groupFieldsMap

        val schema:Schema = if (aggregateInfo.extract_fields.isEmpty) {
            kc.schemaChannel.getSchema(dependencies.head.getSource())
        } else {
            val schemaList:List[KhaosStructField] = schemaMap.map { optInfo =>
                KhaosStructField(optInfo._1, optInfo._2)
            }.toList
            new Schema(schemaList)
        }
        schema
    }

}
